from os.path import join, abspath
import numpy as np
from typing import Set
from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsChoiceMakingTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.framework.shots_prompt.data import ShotsGenerator
from backend.experiment.tokenizer import tokenizer
from backend.experiment.model import context_limit
from backend.experiment.clue.chid.scheme1.implement import CHIDImplement as S1


class CHIDStaticShotsGenerator(ShotsGenerator):
    def __init__(self, implement: 'CHIDImplement'):
        assert isinstance(implement, CHIDImplement)
        self.__implement = implement

    def get_shot(
            self, limit_remain: int, selected_shots: Set[int],
            **kwargs
    ) -> (np.ndarray, int) or None:
        candidate = kwargs['candidate']
        assert candidate in self.__implement.shuffle_idiom_samples.keys()
        examples = self.__implement.shuffle_idiom_samples[candidate]
        next_select = len(selected_shots) - 1
        while next_select < len(examples):
            if next_select not in selected_shots:
                if len(examples[next_select]) <= limit_remain:
                    return examples[next_select], next_select
            next_select += 1
        # 找不到长度比要求更小的了
        return None


class CHIDImplement(S1):
    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsChoiceMakingTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CHIDStaticShotsGenerator(implement=self),
            first_prompt=tokenizer.tokenize('以下句子中成语的使用是正确的：')
        )

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
